package com.algorithm.example.tree;

import java.util.Arrays;

/**
 * <p>
 * B树
 * </p>
 *
 * @author lzt
 * @version 1.0
 * @since 2024/2/26 19:34
 */
public class BTree {
    static class Node {
        /**
         * 关键字
         */
        int[] keys;

        /**
         * 孩子
         */
        Node[] children;

        /**
         * 有效关键字数目
         */
        int keyNumber;

        /**
         * 是否为叶子节点
         */
        boolean leaf = true;

        /**
         * 最小度数（最小孩子数）
         */
        int t;

        public Node(int t) {
            this.t = t;
            this.keys = new int[t * 2 - 1];
            this.children = new Node[t * 2];
        }

        /**
         * 多路查找
         */
        Node get(int key) {
            int a = 0;
            while (a < keyNumber) {
                if (keys[a] == key) {
                    return this;
                }
                if (keys[a] > key) {
                    break;
                }
                a++;
            }

            //叶子节点
            if (leaf) {
                return null;
            }

            // 非叶子节点
            return children[a].get(key);

        }

        /**
         * 移除指定index处的key
         */
        int removerKey(int index) {
            int t = keys[index];
            System.arraycopy(keys, index + 1, keys, index, --keyNumber - index);
            return t;
        }

        /**
         * 移除最左边的key
         */
        int removerLeftmostKey() {
            return removerKey(0);
        }

        /**
         * 移除最右边的key
         */
        int removerRightmostKey() {
            return removerKey(keyNumber - 1);
        }

        /**
         * 移除指定index的child
         */
        Node removerChild(int index) {
            Node child = children[index];
            System.arraycopy(children, index, children, index + 1, keyNumber - index);
            children[keyNumber] = null;
            return child;
        }

        /**
         * 移除最左边的孩子
         */
        Node removerLeftmostChild() {
            return removerChild(0);
        }

        /**
         * 移除最右边的孩子
         */
        Node removerRightmostChild() {
            return removerChild(keyNumber);
        }

        /**
         * index 孩子处左边的兄弟
         */
        Node childLeftSibling(int index) {
            return index > 0 ? children[index - 1] : null;
        }

        /**
         * index 孩子处右边的兄弟
         */
        Node childRightSibling(int index) {
            return index == keyNumber ? null : children[index + 1];
        }

        /**
         * 复制当前节点的所有key和child到target
         */
        void moveToTarget(Node target) {
            int start = target.keyNumber;
            if (!leaf) {
                for (int i = 0; i <= keyNumber; i++) {
                    target.children[start + i] = children[i];
                }
            }
            for (int i = 0; i < keyNumber; i++) {
                target.keys[target.keyNumber++] = keys[i];
            }
        }

        /**
         * 向指定索引插入key
         */
        void insertKey(int key, int index) {
            System.arraycopy(keys, index, keys, index + 1, keyNumber - index);
            keys[index] = key;
            keyNumber++;
        }

        /**
         * 向指定索引插入child
         */
        void insertChild(Node node, int index) {
            System.arraycopy(children, index, children, index + 1, keyNumber - index);
            children[index] = node;
        }

        @Override
        public String toString() {
            return "\nNode{" +
                    "keys=" + Arrays.toString(keys) +
                    ", children=" + Arrays.toString(children) +
                    '}';
//            return Arrays.toString(Arrays.copyOfRange(keys, 0, keyNumber));
        }
    }

    Node root;
    /**
     * 树中节点最小度数
     */
    int t;
    /**
     * 最小key数目
     */
    final int MIN_KEY_NUMBER;
    /**
     * 最大key数目
     */
    final int MAX_KEY_NUMBER;

    public BTree() {
        this(2);
    }

    public BTree(int t) {
        this.t = t;
        root = new Node(t);
        MIN_KEY_NUMBER = t - 1;
        MAX_KEY_NUMBER = 2 * t - 1;
    }


    /**
     * 是否存在
     */
    public boolean contains(int key) {
        return root.get(key) != null;
    }


    /**
     * 新增
     */
    public void put(int key) {
        doPut(root, key, null, 0);
    }

    private void doPut(Node node, int key, Node parent, int index) {
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] == key) {
                return; // 更新
            }
            if (node.keys[i] > key) {
                break; // 找到了插入位置，即为此时的 i
            }
            i++;
        }
        if (node.leaf) {
            node.insertKey(key, i);
        } else {
            doPut(node.children[i], key, node, i);
        }
        if (node.keyNumber == MAX_KEY_NUMBER) {
            split(node, parent, index);
        }
    }

    // 分裂
    private void split(Node left, Node parent, int index) {
        // 判断根节点
        if (parent == null) {
            Node newRoot = new Node(t);
            newRoot.leaf = false;
            newRoot.insertChild(left, 0);
            this.root = newRoot;
            parent = newRoot;
        }

        // 创建right节点，把left中t之后的key和child移动过去
        Node right = new Node(t);
        right.leaf = left.leaf;
        System.arraycopy(left.keys, t, right.keys, 0, t - 1);
        if (!left.leaf) {
            System.arraycopy(left.children, t, right.children, 0, t);
            for (int i = t; i <= left.keyNumber; i++) {
                left.children[i] = null;
            }
        }

        //2.把left的中间key移动到parent中
        left.keyNumber = t - 1;
        right.keyNumber = t - 1;
        int mid = left.keys[t - 1];
        parent.insertKey(mid, index);

        // 3. right 变成parent的index + 1节点
        parent.insertChild(right, index + 1);

    }

    /**
     * 删除
     */
    public void remove(int key) {
        doRemove(null, root, key, 0);
    }

    private void doRemove(Node parent, Node node, int key, int index) {
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] >= key) {
                break;
            }
            i++;
        }

        if (node.leaf) {
            if (!found(node, key, i)) { // case 1 叶子未找到
                return;
            } else {                        // case 2 叶子找到
                node.removerKey(i);
            }
        } else {
            if (!found(node, key, i)) { // case 3 非叶子未找到
                doRemove(node, node.children[i], key, i);
            } else {                        // case 4 非叶子找到
                // 找到后继节点
                Node s = node.children[i + 1];
                while (!s.leaf) {
                    s = s.children[0];
                }
                int skey = s.keys[0];
                // 2.替换待删除的key
                node.keys[i] = skey;
                //3.删除后继节点
                doRemove(node, node.children[i + 1], skey, i + 1);
            }
        }

        if (node.keyNumber < MIN_KEY_NUMBER) {
            // case 5 case 6
            balance(parent, node, index);
        }

    }

    private void balance(Node parent, Node node, int i) {
        if (node == root) {
            if (root.keyNumber == 0 && root.children[0] != null) {
                root = root.children[0];
            }
            return;
        }
        Node left = parent.childLeftSibling(i);
        Node right = parent.childRightSibling(i);
        // 1 左边兄弟富裕，右旋借key
        if (left != null && left.keyNumber > MIN_KEY_NUMBER) {
            node.insertKey(parent.keys[i - 1], 0);
            if (!left.leaf) {
                node.insertChild(left.removerRightmostChild(), 0);
            }
            parent.keys[i - 1] = left.removerRightmostKey();
            return;
        }


        // 2. 右边兄弟富裕，左旋借key
        if (right != null && right.keyNumber > MIN_KEY_NUMBER) {
            node.insertKey(parent.keys[i], node.keyNumber);
            if (!right.leaf) {
                node.insertChild(right.removerLeftmostChild(), node.keyNumber  );
            }
            parent.keys[i] = right.removerLeftmostKey();
            return;
        }

        // 两边都不足，向左合并
        if (left != null) {
            //有左兄弟
            parent.removerChild(i);
            left.insertKey(parent.removerKey(i - 1), left.keyNumber);
            node.moveToTarget(left);
        } else {
            // 没有左兄弟
            parent.removerChild(i + 1);
            node.insertKey(parent.removerKey(i), node.keyNumber);
            right.moveToTarget(node);
        }

    }

    private boolean found(Node node, int key, int i) {
        return i < node.keyNumber && node.keys[i] == key;
    }


    @Override
    public String toString() {
        return "BTree{" +
                "root=" + root +
                '}';
    }
}
